import numpy as np
from scipy.signal import savgol_filter, medfilt
from matplotlib import pyplot as plt

from input import input_data


def savitzky_golay_filter(data:np.ndarray, window_length=15, polyorder=2, mode="nearest", deriv=0):
    '''
    SG滤波
    :param data: y值数组，np.ndarray
    :param window_length: 窗口长度，int，window_length的值越小，曲线越贴近真实曲线；window_length值越大，平滑效果越厉害（备注：该值必须为正奇整数）
    :param polyorder: polyorder为对窗口内的数据点进行k阶多项式拟合，k的值需要小于window_length。，int，k值越大，曲线越贴近真实曲线；k值越小，曲线平滑越厉害。另外，当k值较大时，受窗口长度限制，拟合会出现问题，高频曲线会变成直线。
    :param mode: 确定了要应用滤波器的填充信号的扩展类型。
    :return: 平滑后的矩阵，np.ndarray
    '''
    y = savgol_filter(data,window_length,polyorder,mode=mode,deriv=deriv)
    return y


def median_filter(data: np.ndarray, kernel_size=5) -> np.ndarray:
    '''
    中值滤波
    :param data: 光强数据矩阵，支持多条输入，np.ndarray
    :param kernel_size: kernel_size必须是奇数
    :return: 平滑后的矩阵，np.ndarray
    '''
    filtered_data = medfilt(data, kernel_size = kernel_size)
    return filtered_data

if __name__ == "__main__":
    data,_ = input_data()
    plt.plot( data[0], 'r')
    y = savitzky_golay_filter(data, 15, 2, mode="nearest")
    # 可视化图线
    plt.plot( y[0], 'b', label='savgol')
    # 显示曲线
    plt.show()